import cyaron
from .base import Gener,get_real

CyaronGraphDict = {
    'graph': cyaron.Graph.graph,
    'tree': cyaron.Graph.tree,
    'dag': cyaron.Graph.DAG,
    'hack_spfa': cyaron.Graph.hack_spfa,
    'udag': cyaron.Graph.UDAG,
    'chain': cyaron.Graph.chain,
    'flower': cyaron.Graph.flower,
    'binary_tree': cyaron.Graph.binary_tree
}


class GraphGen(Gener):
    def __init__(self, type_, *args, shuffle=True, **kwargs):
        self.inited = False
        self.type_ = type_
        self.args = args
        self.kwargs = kwargs
        self.graph = None
        self.shuffle = shuffle

    def write(self, file):
        if not self.inited:
            self.inited = True
            args = [get_real(x) for x in self.args]
            kwargs = {}
            for key in self.kwargs:
                kwargs[key] = get_real(self.kwargs[key])
            self.graph = CyaronGraphDict[self.type_](*args, **kwargs)
        file.write(self.graph.to_str(shuffle=self.shuffle))

    def clear(self):
        self.inited = False
